import asyncio
import json
import os
import h5py
import websockets
import requests
from config import bn, db
from app import users
from ws import order
from jose import JWTError, jwt
from method import Bn, Log
from starlette.websockets import WebSocketState
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Cookie

app = APIRouter()

# 从 Binance 订阅推送 Kline
async def binance_kline_stream(websocket: WebSocket, symbol: str, interval: str):
    url = f"{bn.WSbaseUrl}/ws/{symbol}@kline_{interval}"
    try:
        async with websockets.connect(url) as ws:
            async for msg in ws:
                data = json.loads(msg)
                kline = data['k']
                try:
                    # 向所有前端 WebSocket 客户端广播数据
                    await websocket.send_json({
                        "Open": kline["o"],
                        "Close": kline["c"],
                        "High": kline["h"],
                        "Low": kline["l"],
                        "Volume": kline["v"],
                        "OpenTime": kline["t"],
                        "CloseTime": kline["T"]
                    })
                except RuntimeError as e:
                    print(f"Send error: {e}")
                    break  # 终止 Binance 循环
                except Exception as e:
                    print(f"Unhandled send error: {e}")
                    break   
    except Exception as e:
        Log.set_log(f"binance_kline_stream error: {e}", 'error')
        await websocket.close(code=1000, reason=f"Binance socket error:{e}")

@app.websocket("/kline")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    task = None

    try:
        while True:
            data = await websocket.receive_text()
            msg = json.loads(data)

            if msg.get("action") == "subscribe":
                symbol = msg.get("symbol", "btcusdt").lower()
                interval = msg.get("interval", "1m")

                if task:
                    task.cancel()
                task = asyncio.create_task(binance_kline_stream(websocket, symbol, interval))
    except WebSocketDisconnect:
        print("客户端断开")
        if task:
            task.cancel()


@app.websocket("/ticker")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    task = None

    async def binance_ticker_arr():
        try:
            while True:
                await websocket.send_json(order.ticker_data)
                await asyncio.sleep(1)
        except Exception as e:
            if not websocket.client_state.name == "DISCONNECTED":
                await websocket.close(code=1011, reason=f"Binance socket error:{e}")
    
    task = asyncio.create_task(binance_ticker_arr())
    try:
        await task
    except WebSocketDisconnect:
        print("客户端断开")
        if task:
            task.cancel()
    except Exception as e:
        print(f"Unhandled error: {e}")
        if task:
            task.cancel()

# ✅ 定期续期 listenKey（60 分钟过期，建议每 30 分钟刷新）
def keep_alive_listen_key(listen_key, API_KEY):
    headers = {"X-MBX-APIKEY": API_KEY}
    response = requests.put(bn.listenKey, headers=headers, params={"listenKey": listen_key})
    print("🔁 keep-alive status:", response.status_code, response.text)

# ✅ 监听 WebSocket 推送
async def listen_account_data(websocket: WebSocket, listen_key):
    ws_url = f"{bn.WSbaseUrl}/ws/{listen_key}"
    try:
        async with websockets.connect(ws_url) as ws:
            async for msg in ws:
                data = json.loads(msg)
                await websocket.send_json(data)
    except Exception as e:
        Log.set_log(f"listen_account_data error: {e}", 'error')
        if websocket.application_state == WebSocketState.CONNECTED:
            await websocket.close(code=1011, reason=f"Binance socket error: {e}")

# ✅ 定期续期 listenKey
async def auto_keep_alive(listen_key, API_KEY):
    while True:
        await asyncio.sleep(30 * 60)  # 每30分钟
        keep_alive_listen_key(listen_key, API_KEY)

@app.websocket("/account")
async def websocket_endpoint(websocket: WebSocket, token: str = Cookie(None)):
    await websocket.accept()

    if token is None:
        Log.set_log(f"websocket_endpoint error: 获取token失败", 'error')
        await websocket.close(code=4001, reason="获取token失败")
        return
    try:
        payload = jwt.decode(token, users.SECRET_KEY, algorithms=[users.ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            await websocket.close(code=4001, reason="无效token")
            return
    except Exception as e:
        Log.set_log(f"websocket_endpoint error: {e}", 'error')
        await websocket.close(code=4001, reason=f"token解析失败: {e}")
        return

    task = None
    binance_task = None
    try:
        while True:
            # 获取前端send数据
            data = await websocket.receive_text()
            msg = json.loads(data)

            if msg.get("action") == "subscribe" and username:
                # 获取listenKey
                my_binance = Bn.Binance(username)
                listen = my_binance.synced(bn.listenKey, {}, 'post')
                listen_key = listen["listenKey"]
                # listenKey定期续期
                task = asyncio.create_task(auto_keep_alive(listen_key, my_binance.headers['X-MBX-APIKEY']))
                # 根据listenKey获取用户币安合约账户信息变化
                binance_task = asyncio.create_task(listen_account_data(websocket, listen_key))
    except WebSocketDisconnect:
        print("客户端断开")
        if task:
            task.cancel()
        if binance_task:
            binance_task.cancel()

